# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Load COCO dataset."""
import os

import numpy as np
from PIL import Image
from pycocotools.coco import COCO

from mindvision.engine.class_factory import ClassFactory, ModuleType


@ClassFactory.register(ModuleType.DATASET)
class COCOYoloDataset:
    """YOLOV4 Dataset for COCO."""

    def __init__(self,
                 root,
                 ann_file,
                 remove_images_without_annotations=True,
                 filter_crowd_anno=True,
                 is_training=True):
        self.coco = COCO(ann_file)
        self.root = root
        self.img_ids = list(sorted(self.coco.imgs.keys()))
        self.filter_crowd_anno = filter_crowd_anno
        self.is_training = is_training

        # filter images without any annotations
        if remove_images_without_annotations:
            img_ids = []
            for img_id in self.img_ids:
                ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=None)
                anno = self.coco.loadAnns(ann_ids)
                if has_valid_annotation(anno):
                    img_ids.append(img_id)
            self.img_ids = img_ids

        self.categories = {
            cat["id"]: cat["name"] for cat in self.coco.cats.values()
        }

        self.cat_ids_to_continuous_ids = {
            v: i for i, v in enumerate(self.coco.getCatIds())
        }
        self.continuous_ids_cat_ids = {
            v: k for k, v in self.cat_ids_to_continuous_ids.items()
        }

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            (img, target) (tuple):
            target is a dictionary contains "bbox", "segmentation" or "keypoints",
            generated by the image's annotation. img is a PIL image.
        """
        coco = self.coco
        img_id = self.img_ids[index]
        img_path = coco.loadImgs(img_id)[0]["file_name"]
        img = Image.open(os.path.join(self.root, img_path)).convert("RGB")
        if not self.is_training:
            return img, img_id

        ann_ids = coco.getAnnIds(imgIds=img_id)
        target = coco.loadAnns(ann_ids)
        # filter crowd annotations
        if self.filter_crowd_anno:
            annos = [anno for anno in target if anno["iscrowd"] == 0]
        else:
            annos = [anno for anno in target]

        target = {}
        boxes = [anno["bbox"] for anno in annos]
        target["bboxes"] = boxes

        classes = [anno["category_id"] for anno in annos]
        classes = [self.cat_ids_to_continuous_ids[cl] for cl in classes]
        target["labels"] = classes

        bboxes = target['bboxes']
        labels = target['labels']
        out_target = []
        for bbox, label in zip(bboxes, labels):
            tmp = []
            # convert to [x_min y_min x_max y_max]
            bbox = xywh2xyxy(bbox)
            tmp.extend(bbox)
            tmp.append(int(label))
            # tmp [x_min y_min x_max y_max, label]
            out_target.append(tmp)
        return img, out_target, [], [], [], [], []

    def __len__(self):
        return len(self.img_ids)


def xywh2xyxy(bbox):
    """xywh convert into xyxy format."""
    x_min = bbox[0]
    y_min = bbox[1]
    w = bbox[2]
    h = bbox[3]
    return [x_min, y_min, x_min + w, y_min + h]


def _has_only_empty_bbox(anno):
    return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno)


def _count_visible_keypoints(anno):
    return sum(sum(1 for v in ann["keypoints"][2::3] if v > 0) for ann in anno)


def has_valid_annotation(anno):
    """Check annotation file."""
    min_keypoints_per_image = 10
    # if it's empty, there is no annotation
    if not anno:
        return False
    # if all boxes have close to zero area, there is no annotation
    if _has_only_empty_bbox(anno):
        return False
    # keypoints task have a slight different criteria for considering
    # if an annotation is valid
    if "keypoints" not in anno[0]:
        return True
    # for keypoint detection tasks, only consider valid images those
    # containing at least min_keypoints_per_image
    if _count_visible_keypoints(anno) >= min_keypoints_per_image:
        return True
    return False


def ann_to_mask(ann, height, width):
    """Convert annotation to RLE and then to binary mask."""
    from pycocotools import mask as maskHelper
    segm = ann['segmentation']
    if isinstance(segm, list):
        rles = maskHelper.frPyObjects(segm, height, width)
        rle = maskHelper.merge(rles)
    elif isinstance(segm['counts'], list):
        rle = maskHelper.frPyObjects(segm, height, width)
    else:
        rle = ann['segmentation']
    m = maskHelper.decode(rle)
    return m


@ClassFactory.register(ModuleType.DATASET)
class COCODataset:
    """COCO Dataset."""

    def __init__(self,
                 root,
                 ann_file,
                 remove_images_without_annotations=False,
                 filter_crowd_anno=True,
                 is_training=True,
                 with_mask=True):
        self.coco = COCO(ann_file)
        self.root = root
        # self.img_ids = list(sorted(self.coco.imgs.keys()))
        self.img_ids = self.coco.getImgIds()
        self.filter_crowd_anno = filter_crowd_anno
        self.is_training = is_training
        self.with_mask = with_mask

        # filter images without any annotations
        if remove_images_without_annotations:
            img_ids = []
            for img_id in self.img_ids:
                ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=None)
                anno = self.coco.loadAnns(ann_ids)
                if has_valid_annotation(anno):
                    img_ids.append(img_id)
            self.img_ids = img_ids

        train_cls = ['background', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
                     'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
                     'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
                     'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
                     'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
                     'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
                     'kite', 'baseball bat', 'baseball glove', 'skateboard',
                     'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
                     'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
                     'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
                     'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
                     'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
                     'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
                     'refrigerator', 'book', 'clock', 'vase', 'scissors',
                     'teddy bear', 'hair drier', 'toothbrush']

        self.train_cls_dict = {}
        for i, cls in enumerate(train_cls):
            self.train_cls_dict[cls] = i

        self.categories = {
            cat["id"]: cat["name"] for cat in self.coco.cats.values()
        }

        self.cat_ids_to_continuous_ids = {
            v: i for i, v in enumerate(self.coco.getCatIds())
        }
        self.continuous_ids_cat_ids = {
            v: k for k, v in self.cat_ids_to_continuous_ids.items()
        }

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            (img, target) (tuple):
            target is a dictionary contains "bbox", "segmentation" or "keypoints",
            generated by the image's annotation. img is a PIL image.
        """
        coco = self.coco
        img_id = self.img_ids[index]
        img_path = coco.loadImgs(img_id)[0]["file_name"]
        image = Image.open(os.path.join(self.root, img_path)).convert("RGB")

        if self.with_mask:
            instance_mask = []
            image_height = coco.imgs[img_id]["height"]
            image_width = coco.imgs[img_id]["width"]

        ann_ids = coco.getAnnIds(imgIds=img_id)
        anno = coco.loadAnns(ann_ids)

        annotation = []
        for label in anno:
            bbox = label["bbox"]
            class_name = self.categories[label["category_id"]]

            if self.filter_crowd_anno and label["iscrowd"] == 1:
                continue

            x1, x2 = bbox[0], bbox[0] + bbox[2]
            y1, y2 = bbox[1], bbox[1] + bbox[3]

            if self.with_mask:
                m = ann_to_mask(label, image_height, image_width)
                instance_mask.append(m)
            annotation.append(
                [x1, y1, x2, y2] + [self.train_cls_dict[class_name]] + [
                    int(label["iscrowd"])])

        if annotation:
            annotation = np.array(annotation, dtype=np.int32)
            if self.with_mask:
                instance_mask = np.stack(instance_mask, axis=0).astype(np.int32)
                mask = np.array(instance_mask)
                mask_shape = np.array(instance_mask.shape, dtype=np.int32)
        else:
            annotation = np.array([[0, 0, 0, 0, 0, 1]])
            if self.with_mask:
                mask = np.zeros([1, image_height, image_width], dtype=np.int32)
                mask_shape = np.array([1, image_height, image_width], dtype=np.int32)
        if self.with_mask:
            return [img_id], image, annotation, mask, mask_shape
        return [img_id], image, annotation

    def __len__(self):
        return len(self.img_ids)


@ClassFactory.register(ModuleType.DATASET)
class COCOSsdDataset:
    """ssd dataset for coco"""
    def __init__(self,
                 root,
                 ann_file,
                 filter_crowd_anno=True,
                 is_training=True):
        self.coco = COCO(ann_file)
        self.root = root
        self.img_ids = list(sorted(self.coco.imgs.keys()))
        # self.img_ids = self.coco.getImgIds()
        self.filter_crowd_anno = filter_crowd_anno
        self.is_training = is_training

        train_cls = ['background', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
                     'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
                     'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
                     'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
                     'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
                     'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
                     'kite', 'baseball bat', 'baseball glove', 'skateboard',
                     'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
                     'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
                     'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
                     'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
                     'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
                     'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
                     'refrigerator', 'book', 'clock', 'vase', 'scissors',
                     'teddy bear', 'hair drier', 'toothbrush']
        self.train_cls_dict = {}
        for i, cls in enumerate(train_cls):
            self.train_cls_dict[cls] = i

        self.categories = {
            cat["id"]: cat["name"] for cat in self.coco.cats.values()
        }

        img_ids = []
        image_path_dict = {}
        image_anno_dict = {}

        for img_id in self.img_ids:
            image_info = self.coco.loadImgs(img_id)
            file_name = image_info[0]["file_name"]
            anno_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=None)
            anno = self.coco.loadAnns(anno_ids)
            # image_path = os.path.join(self.root, "train2017", file_name)
            image_path = os.path.join(self.root, file_name)
            annos = []
            iscrowd = False
            for label in anno:
                bbox = label["bbox"]
                class_name = self.categories[label["category_id"]]  # key error
                iscrowd = iscrowd or label["iscrowd"]
                if class_name in train_cls:
                    x_min, x_max = bbox[0], bbox[0] + bbox[2]
                    y_min, y_max = bbox[1], bbox[1] + bbox[3]
                    annos.append(list(map(round, [y_min, x_min, y_max, x_max])) + [self.train_cls_dict[class_name]])

            if not self.is_training and iscrowd:
                continue
            if len(annos) >= 1:
                img_ids.append(img_id)
                image_path_dict[img_id] = image_path
                image_anno_dict[img_id] = np.array(annos)

        self.img_ids = img_ids
        self.image_path_dict = image_path_dict
        self.image_anno_dict = image_anno_dict

        self.cat_ids_to_continuous_ids = {
            v: i for i, v in enumerate(self.coco.getCatIds())
        }
        self.continuous_ids_cat_ids = {
            v: k for k, v in self.cat_ids_to_continuous_ids.items()
        }

    def __getitem__(self, index):
        coco = self.coco
        img_id = self.img_ids[index]
        img_path = coco.loadImgs(img_id)[0]["file_name"]
        image = Image.open(os.path.join(self.root, img_path)).convert("RGB")
        return img_id, image, self.image_anno_dict[img_id]

    def __len__(self):
        return len(self.img_ids)
